import networkx as nx
import torch
import torch.nn as nn
import numpy as np
import inspect
from torch.fx import symbolic_trace, GraphModule
from typing import Dict, List, Tuple, Any, Optional, Union


class Graph:
    """自定义图结构，基于 NetworkX 存储模型信息"""

    def __init__(self, name: str = ""):
        self.name = name
        self.graph = nx.DiGraph()
        self.input_nodes: List[str] = []
        self.output_nodes: List[str] = []
        self.node_attrs: Dict[str, Dict] = {}
        self.param_map: Dict[str, str] = {}  # 参数名到节点名的映射

    def add_node(self, name: str, **attrs):
        """添加节点"""
        self.graph.add_node(name)
        self.node_attrs[name] = attrs

    def add_edge(self, source: str, target: str):
        """添加边"""
        self.graph.add_edge(source, target)

    def topological_sort(self) -> List[str]:
        """拓扑排序"""
        return list(nx.topological_sort(self.graph))


def get_module_metadata(module: nn.Module, module_path: str) -> Dict[str, Any]:
    """获取模块的元数据，包括初始化参数、参数和缓冲区"""
    metadata = {
        "init_args": {},
        "parameters": {},
        "buffers": {},
        "module_path": module_path,  # 添加模块路径信息
    }

    # 获取模块类的__init__方法签名
    init_signature = inspect.signature(module.__class__.__init__)

    # 提取初始化参数
    for param_name, param in init_signature.parameters.items():
        if param_name == "self":
            continue

        # 尝试从模块实例中获取参数值
        if hasattr(module, param_name):
            value = getattr(module, param_name)
            # 只存储基本类型和元组/列表
            if isinstance(value, (int, float, bool, str, tuple, list)):
                metadata["init_args"][param_name] = value
            elif isinstance(value, torch.Tensor):
                # 存储张量的形状和类型
                metadata["init_args"][param_name] = {
                    "shape": tuple(value.shape),
                    "dtype": str(value.dtype),
                }

    # 提取参数 - 使用原始路径名称
    for name, param in module.named_parameters(recurse=False):
        # 使用原始路径名称，不添加额外前缀
        metadata["parameters"][f"{module_path}.{name}"] = param.detach().cpu().numpy()

    # 提取缓冲区
    for name, buffer in module.named_buffers(recurse=False):
        # 使用原始路径名称
        metadata["buffers"][f"{module_path}.{name}"] = buffer.detach().cpu().numpy()

    return metadata


def model_to_graph(model: nn.Module, input_shape: Tuple[int, ...]) -> Graph:
    """
    将 PyTorch 模型转换为 NetworkX 图
    Args:
        model: PyTorch 模型
        input_shape: 输入张量形状 (用于追踪模型)
    Returns:
        自定义图结构
    """
    # 使用 FX 追踪模型
    graph_module = symbolic_trace(model)

    # 创建自定义图
    graph = Graph(name=model.__class__.__name__)
    node_map = {}  # 映射 FX 节点名到自定义节点名

    # 提取所有模块信息
    modules = dict(model.named_modules())

    # 添加输入节点
    input_names = []
    for i, node in enumerate(graph_module.graph.nodes):
        if node.op == "placeholder":
            node_name = f"input_{i}"
            graph.add_node(node_name, type="Input", shape=input_shape, op_class=None)
            graph.input_nodes.append(node_name)
            input_names.append(node_name)
            node_map[node.name] = node_name

    # 处理计算节点 - 使用原始节点名称
    for node in graph_module.graph.nodes:
        if node.op in ["call_module", "call_function", "call_method"]:
            # 使用原始节点名称（如 'layer1_0_conv1'）作为节点标识
            node_name = node.name

            # 提取操作信息
            if node.op == "call_module":
                module = modules[node.target]
                op_class = type(module)

                # 获取模块元数据 - 传递完整模块路径
                metadata = get_module_metadata(module, node.target)

                # 添加节点
                graph.add_node(node_name, type=node.op, op_class=op_class, **metadata)

                # 记录参数到节点的映射
                for param_name in metadata["parameters"]:
                    graph.param_map[param_name] = node_name
                for buffer_name in metadata["buffers"]:
                    graph.param_map[buffer_name] = node_name

            elif node.op == "call_function":
                # 处理函数调用
                graph.add_node(
                    node_name,
                    type=node.op,
                    op_class=node.target,
                    init_args={},
                    parameters={},
                    buffers={},
                    module_path="",
                )

            elif node.op == "call_method":
                # 处理方法调用
                graph.add_node(
                    node_name,
                    type=node.op,
                    op_class=getattr(torch.Tensor, node.target),
                    init_args={},
                    parameters={},
                    buffers={},
                    module_path="",
                )

            node_map[node.name] = node_name

            # 添加边
            for arg in node.args:
                if isinstance(arg, torch.fx.Node) and arg.name in node_map:
                    graph.add_edge(node_map[arg.name], node_name)

    # 添加输出节点
    for node in graph_module.graph.nodes:
        if node.op == "output":
            output_name = f"output_{len(graph.output_nodes)}"
            graph.add_node(output_name, type="Output", op_class=None, module_path="")
            graph.output_nodes.append(output_name)

            # 连接到输出节点
            for arg in node.args:
                if isinstance(arg, torch.fx.Node) and arg.name in node_map:
                    graph.add_edge(node_map[arg.name], output_name)

    return graph


def graph_to_model(graph: Graph) -> nn.Module:
    """
    将 NetworkX 图转换回 PyTorch 模型
    Args:
        graph: 自定义图结构
    Returns:
        PyTorch 模型
    """

    class ReconstructedModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.modules_dict = nn.ModuleDict()
            self.node_order = graph.topological_sort()

            # 创建所有模块并立即设置参数
            for node_name in self.node_order:
                attrs = graph.node_attrs[node_name]
                node_type = attrs["type"]

                if node_type == "call_module":
                    op_class = attrs["op_class"]
                    init_args = attrs.get("init_args", {})

                    # 确保卷积层的bias设置正确
                    if op_class in [nn.Conv1d, nn.Conv2d, nn.Conv3d]:
                        # 检查原始模型是否有bias参数
                        has_bias = any("bias" in param for param in attrs["parameters"])
                        init_args["bias"] = has_bias

                    # 创建模块实例
                    module = op_class(**init_args)

                    # 立即设置参数和缓冲区
                    self._set_module_parameters(
                        module, attrs, graph.param_map, node_name
                    )

                    # 添加到模块字典
                    self.add_module(node_name, module)
                    self.modules_dict[node_name] = module

        def _set_module_parameters(
            self, module: nn.Module, attrs: Dict, param_map: Dict, node_name: str
        ):
            """为模块设置参数和缓冲区"""
            # 获取模块的当前参数和缓冲区
            module_params = dict(module.named_parameters())
            module_buffers = dict(module.named_buffers())

            # 设置参数
            for param_name, array in attrs.get("parameters", {}).items():
                # 只处理属于当前节点的参数
                if param_map.get(param_name) == node_name:
                    base_name = param_name.split(".")[-1]
                    tensor = torch.from_numpy(array)

                    # 如果模块有该参数，则设置值
                    if base_name in module_params:
                        if module_params[base_name].shape == tensor.shape:
                            module_params[base_name].data.copy_(tensor)
                    else:
                        # 注册新参数（如果需要）
                        param = nn.Parameter(tensor)
                        module.register_parameter(base_name, param)

            # 设置缓冲区
            for buffer_name, array in attrs.get("buffers", {}).items():
                if param_map.get(buffer_name) == node_name:
                    base_name = buffer_name.split(".")[-1]
                    tensor = torch.from_numpy(array)

                    # 如果模块有该缓冲区，则设置值
                    if base_name in module_buffers:
                        if module_buffers[base_name].shape == tensor.shape:
                            module_buffers[base_name].data.copy_(tensor)
                    else:
                        # 注册新缓冲区
                        module.register_buffer(base_name, tensor)

        def forward(self, x):
            node_results = {}

            # 处理输入节点
            for i, input_name in enumerate(graph.input_nodes):
                # 对于多个输入的情况
                if isinstance(x, (list, tuple)) and i < len(x):
                    node_results[input_name] = x[i]
                else:
                    node_results[input_name] = x

            # 按拓扑顺序处理节点
            for node_name in self.node_order:
                attrs = graph.node_attrs[node_name]
                node_type = attrs["type"]

                # 跳过输入和输出节点（它们不执行计算）
                if node_type in ["Input", "Output"]:
                    continue

                # 获取所有输入
                inputs = []
                for predecessor in graph.graph.predecessors(node_name):
                    if predecessor in node_results:
                        inputs.append(node_results[predecessor])

                if node_type == "call_module":
                    # 执行模块
                    module = self.modules_dict[node_name]

                    # 处理多个输入
                    if len(inputs) > 1:
                        output = module(*inputs)
                    elif len(inputs) == 1:
                        output = module(inputs[0])
                    else:
                        output = module()

                    node_results[node_name] = output

                elif node_type == "call_function":
                    # 执行函数
                    func = attrs["op_class"]
                    node_results[node_name] = func(*inputs)

                elif node_type == "call_method":
                    # 执行方法 - 假设第一个输入是张量
                    method = attrs["op_class"]
                    if inputs:
                        tensor_input = inputs[0]
                        other_inputs = inputs[1:]
                        node_results[node_name] = method(tensor_input, *other_inputs)

            # 处理输出节点
            outputs = []
            for output_name in graph.output_nodes:
                # 获取输出节点的所有输入
                for predecessor in graph.graph.predecessors(output_name):
                    if predecessor in node_results:
                        outputs.append(node_results[predecessor])

            # 根据输出数量返回适当的结果
            if len(outputs) == 1:
                return outputs[0]
            elif len(outputs) > 1:
                return tuple(outputs)
            else:
                return None

    return ReconstructedModel()


# 测试函数
def test_resnet18_conversion():
    import torchvision.models as models

    # 创建 ResNet18 模型
    resnet18 = models.resnet18(pretrained=True)

    # 转换为图
    graph = model_to_graph(resnet18, input_shape=(1, 3, 224, 224))
    print(f"Graph created with {len(graph.graph.nodes)} nodes")

    # 转换回模型
    reconstructed_model = graph_to_model(graph)

    # 验证模型结构
    input_tensor = torch.randn(1, 3, 224, 224)

    # 原始模型输出
    resnet18.eval()
    reconstructed_model.eval()

    with torch.no_grad():
        original_output = resnet18(input_tensor)
        reconstructed_output = reconstructed_model(input_tensor)

    # 检查输出形状是否匹配
    print(f"Original output shape: {original_output.shape}")
    print(f"Reconstructed output shape: {reconstructed_output.shape}")

    # 检查参数是否匹配
    original_params = dict(resnet18.named_parameters())
    reconstructed_params = dict(reconstructed_model.named_parameters())

    # 检查输出值是否接近
    diff = torch.abs(original_output - reconstructed_output).max().item()
    print(f"Max output difference: {diff}")

    torch.onnx.export(
        resnet18,
        input_tensor,
        "original.onnx",
        opset_version=13,  # 使用支持量化算子的opset
        input_names=["input"],
        output_names=["output"],
        do_constant_folding=False,  # 禁用常量折叠
        verbose=True,
    )

    torch.onnx.export(
        reconstructed_model,
        input_tensor,
        "reconstructed.onnx",
        opset_version=13,  # 使用支持量化算子的opset
        input_names=["input"],
        output_names=["output"],
        do_constant_folding=False,  # 禁用常量折叠
        verbose=True,
    )


if __name__ == "__main__":
    test_resnet18_conversion()
